package com.googlecode.garbagecan.dbtest.vertica;

import com.googlecode.garbagecan.dbtest.Callback;
import com.googlecode.garbagecan.dbtest.util.DBUtil;
import org.apache.commons.dbutils.QueryRunner;
import org.apache.commons.dbutils.ResultSetHandler;

import java.io.BufferedReader;
import java.io.ByteArrayInputStream;
import java.io.IOException;
import java.io.InputStreamReader;
import java.nio.Buffer;
import java.sql.*;
import java.util.ArrayList;
import java.util.Formatter;
import java.util.List;
import java.util.Locale;

// mvn install:install-file -DgroupId=com.vertica -DartifactId=vertica -Dversion=4.1.7_jdk_5 -Dpackaging=jar -Dfile=vertica_4.1.7_jdk_5.jar -DgeneratePom=true
public class TestStorage {
	private static final String driver = "com.vertica.Driver";

	// Vertica 4x
	private static final String url_4x = "jdbc:vertica://perfqa04:5550/fkong";
	private static final String username_4x = "pa8";
	private static final String password_4x = "fkong";

	// Vertica 5x
	private static final String url_5x = "jdbc:vertica://perf2:5433/pa8_sim";
	private static final String username_5x = "pa8";
	private static final String password_5x  = "pa8_sim";

	// Vertica 5x
//	private static final String url_5x = "jdbc:vertica://172.17.23.121:5550/fkong";
//	private static final String username_5x = "dbadmin";
//	private static final String password_5x  = "fkong";

	public static void main(String[] args) throws Exception {
		TestStorage ts = new TestStorage();
		ts.test();
		//ts.getLicenseInfo();
	}

	private void test() throws Exception {
		final Connection conn4x = get4XConnection();
		doIt(new Callback() {
			public void execute() throws Exception {
				long size = getTablesSize(conn4x, username_4x);
				System.out.println("-----------------------total size: " + size);
				System.out.println("\nTotal Size: " + getFormattedSize(size));
			}
		});
		conn4x.close();

		final Connection conn5x = get5XConnection();
		doIt(new Callback() {
			public void execute() throws Exception {
				long size = getTablesSize(conn5x, username_5x);
				System.out.println("-----------------------total size: " + size);
				System.out.println("\nTotal Size: " + getFormattedSize(size));
			}
		});
		conn5x.close();
	}

	private String getLicenseInfo(Connection conn) throws Exception {
		String sql = "SELECT DISPLAY_LICENSE();";
		QueryRunner queryRunner = new QueryRunner();
		String license = queryRunner.query(conn, sql, new ResultSetHandler<String>() {
			public String handle(ResultSet rs) throws SQLException {
				if (!rs.next()) {
					return null;
				}
				String license = rs.getString(1);
				try {
					ByteArrayInputStream bais = new ByteArrayInputStream(license.getBytes());
					BufferedReader br = new BufferedReader(new InputStreamReader(bais));
					String line = null;
					while((line = br.readLine()) != null) {
						System.out.println(line);
					}
					bais.close();
				} catch(IOException ex) {
					ex.printStackTrace();
				}
				//System.out.println(license);
				return license;
			}
		});
		return license;
	}

	private long getTablesSize(Connection conn, String username) throws Exception {
		long size = 0;
		List<String> tables = getRawTables(conn, username);
		for (final String table : tables) {
			size += getTableSize(conn, table);
		}
		return size;
	}

	private long getTableSize(Connection conn, final String table) throws Exception {
		String sTotalRecordNum = DBUtil.getSingleValue(conn, "SELECT count(*) FROM " + table);
		sTotalRecordNum = (sTotalRecordNum == null) ? "0" : sTotalRecordNum;

		String sql = "SELECT * FROM " + table + " where 1=0";
		PreparedStatement ps = conn.prepareStatement(sql);
		ResultSet rs = ps.executeQuery();
		ResultSetMetaData rsmd = rs.getMetaData();
		StringBuffer sbSQL = new StringBuffer();
		sbSQL.append("SELECT ");
		int colCount = rsmd.getColumnCount();
		for (int i = 1; i <= colCount; i++) {
			String columnName = rsmd.getColumnName(i);
			String columnType = rsmd.getColumnTypeName(i);
			if (i != 1) {
				sbSQL.append("+");
			}
			if (columnType.equalsIgnoreCase("binary") || columnType.equalsIgnoreCase("varbinary")) {
				sbSQL.append("SUM(LENGTH(NVL(").append(columnName).append(", '')))");
			} else {
				sbSQL.append("SUM(LENGTH(NVL(").append(columnName).append("::VARCHAR, '')))");
			}
		}
		sbSQL.append(" FROM ").append(table);
		rs.close();
		ps.close();

		String sTotalRecordSize = DBUtil.getSingleValue(conn, sbSQL.toString());
		long totalRecordSize = (sTotalRecordSize == null) ? 0 : Long.parseLong(sTotalRecordSize);
		sTotalRecordSize = getFormattedSize(totalRecordSize);

		System.out.println("==================================================");
		System.out.println("Table               " + table);
		System.out.println("--------------------------------------------------");
		System.out.println("Total Record Number " + sTotalRecordNum);
		System.out.println("Total Record Size   " + sTotalRecordSize);
		System.out.println("Execute SQL         " + sbSQL.toString());
		System.out.println("--------------------------------------------------");
		
		return totalRecordSize;
	}

	private List<String> getRawTables(Connection conn, String username) throws Exception {
		String sql = "SELECT TABLE_SCHEMA||'.'||TABLE_NAME AS TABLE FROM TABLES "
			+ " WHERE OWNER_NAME='"+username+"' "
			+ " ORDER BY TABLE_SCHEMA, TABLE_NAME";
		QueryRunner queryRunner = new QueryRunner();
		List<String> tables = queryRunner.query(conn, sql, new ResultSetHandler<List<String>>() {
			public List<String> handle(ResultSet rs) throws SQLException {
				List<String> tables = new ArrayList<String>();
				while (rs.next()) {
					String table = rs.getString(1);
					tables.add(table);
				}
				return tables;
			}
		});
		return tables;
	}

	private String getFormattedSize(long size) {
		String ret = "";
		Formatter formatter = new Formatter();
		float fSize = (float)size;
		if (fSize >= ((1L<<40)) ) {
			ret = formatter.format("%.2f", fSize / 1024 / 1024 / 1024 / 1024).toString() + " TB";
		} else if (fSize >= (1L<<30) ) {
			ret = formatter.format("%.2f", fSize / 1024 / 1024 / 1024).toString() + " GB";
		} else if (fSize >= (1L<<20) ) {
			ret = formatter.format("%.2f", fSize / 1024 / 1024).toString() + " MB";
		} else if (fSize >= (1L<<10) ) {
			ret = formatter.format("%.2f", fSize / 1024 ).toString() + " KB";
		} else {
			ret = formatter.format("%.2f", fSize).toString() + " B";
		}
		return ret;
	}

	private Connection get4XConnection() throws Exception {
		Class.forName(driver);
		Connection conn = DriverManager.getConnection(url_4x, username_4x, password_4x);
		conn.setAutoCommit(false);
		return conn;
	}

	private Connection get5XConnection() throws Exception {
		Class.forName(driver);
		Connection conn = DriverManager.getConnection(url_5x, username_5x, password_5x);
		conn.setAutoCommit(false);
		return conn;
	}
	
	private void doIt(Callback callback) throws Exception {
		long start = System.currentTimeMillis();
		callback.execute();
		long stop = System.currentTimeMillis();
		System.out.println("Time used: " + (stop - start) + " ms");
	}
}
